import yaml
import os
from logger import Logger

class Config:
    def __init__(self, config_path='config.yml'):
        self.logger = Logger(name='config')
        self.config_path = config_path
        self.logger.info(f"加载配置文件: {config_path}")
        self.config = self._load_config()
        
    def _load_config(self):
        """加载配置文件"""
        if not os.path.exists(self.config_path):
            self.logger.warning(f"配置文件不存在，创建默认配置: {self.config_path}")
            return self._create_default_config()
        
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            self.logger.info("配置加载成功")
            return config
        except Exception as e:
            self.logger.error(f"加载配置文件失败: {str(e)}")
            raise
    
    def _create_default_config(self):
        """创建默认配置"""
        config = {
            'data': {
                'data_dir': 'dataset',
                'img_size': 224,
                'batch_size': 32
            },
            'model': {
                'type': 'transfer',  # 'basic' or 'transfer'
                'save_dir': 'models',
                'learning_rate': 0.001
            },
            'training': {
                'epochs': 10,
                'validation_split': 0.2,
                'early_stopping_patience': 5
            }
        }
        
        # 保存默认配置
        self.save_config(config)
        return config
    
    def save_config(self, config):
        """保存配置到文件"""
        with open(self.config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False)
    
    def get(self, key, default=None):
        """获取配置值"""
        keys = key.split('.')
        value = self.config
        for k in keys:
            if isinstance(value, dict):
                value = value.get(k)
            else:
                return default
        return value if value is not None else default 